#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Datetime: 2020/11/17 20:24
# @Author  : CHEN Wang
# @Site    : 
# @File    : bond_price_related.py
# @Software: PyCharm 

"""
脚本说明：获取债券行情数据相关函数
"""

import pandas as pd
from quant_researcher.quant.project_tool.celebrity import get_lst_condition
from quant_researcher.quant.project_tool.db_operator import db_conn, my_mysql
from quant_researcher.quant.project_tool import time_tool

# 上清所债券估值表
BOND_HQ_TABLE = 'bnd_dr_shclearvaluation'
# 可转债行情表
CONVERTIBLE_BOND_HQ_TABLE = 'bnd_conv_daily_quote'
# 债券编码信息表
BOND_INFO_TABLE = 'bnd_bd_bondcode'


def get_bond_quote(bond_code, start_date, end_date, select=None, **kwargs):
    """
    得到债券的行情数据
    表非常大，暂时不支持only_latest

    :param str or list bond_code: 需要的债券代码列表，债券数据太多，必须传入该参数
    :param str or int or datetime.date or datetime.datetime or int start_date: 开始时间，支持带杠字符串、不带杠字符串、datetime、date、int格式
    :param str or int or datetime.date or datetime.datetime or int end_date: 结束时间，支持带杠字符串、不带杠字符串、datetime、date、int格式
    :param list select: 需要哪几列，默认所有
    :param kwargs:
        - conn: 数据库连接
    :return: DataFrame
                - bond_code: 债券代码
                - end_date: 截止时间
                - eval_yield: 估价收益率(%)
                - eval_aduration: 估价修正久期
                - deeval_fullprice: 日终估价全价(元)
    """
    conn = kwargs.pop('conn', None)
    tmp_conn = db_conn.get_basic_data_conn() if conn is None else conn
    if select is None:
        select = ['bond_code', 'end_date', 'eval_yield', 'eval_aduration', 'deeval_fullprice']

    if isinstance(bond_code, str):
        bond_code = [bond_code]
    bond_where_condition = get_lst_condition(bond_code, 'bond_code')

    start_date = time_tool.format_date_str(start_date, '%Y%m%d')
    end_date = time_tool.format_date_str(end_date, '%Y%m%d')

    bond_hq = my_mysql.read_v2(select=select, sfrom=BOND_HQ_TABLE,
                               where=[f"end_date >= '{start_date}'", f"end_date <= '{end_date}'", bond_where_condition],
                               conn=tmp_conn)
    if bond_hq.empty:
        bond_hq = pd.DataFrame(columns=select)

    # 可转债行情表中没有bond_code字段，需要得到bond_code和secinner_code的对应关系
    # 可转债估值指标和债券估值指标字段对应字典
    indicator_map_dict = {'eval_yield': 'ytm_cl',
                          'eval_aduration': 'modified_duration_cl',
                          'deeval_fullprice': 'close',
                          'bond_code': 'secinner_code',
                          'end_date': 'trade_date'}
    select_conv = [indicator_map_dict[x] if x in indicator_map_dict.keys() else x for x in select]

    conv_bond_code = my_mysql.read_v2(select=['secinner_code', 'bond_code'], sfrom=BOND_INFO_TABLE,
                                      where=bond_where_condition, conn=tmp_conn)
    if conv_bond_code.empty:
        conv_bond_hq = pd.DataFrame(columns=select)
    else:
        secinner_code_where_condition = get_lst_condition(conv_bond_code['secinner_code'].tolist(), 'secinner_code')
        conv_bond_hq = my_mysql.read_v2(select=select_conv, sfrom=CONVERTIBLE_BOND_HQ_TABLE,
                                        where=[f"trade_date >= '{start_date}'", f"trade_date <= '{end_date}'", secinner_code_where_condition],
                                        conn=tmp_conn)
        conv_bond_hq.columns = select
        conv_bond_hq = conv_bond_hq.rename(columns={'bond_code': 'secinner_code'})
        conv_bond_hq = conv_bond_hq.merge(conv_bond_code, how='left', on='secinner_code')
        conv_bond_hq = conv_bond_hq.drop(columns=['secinner_code'])
    result = pd.concat([bond_hq, conv_bond_hq])
    if conn is None:
        tmp_conn.close()

    result['end_date'] = result['end_date'].apply(time_tool.format_date_str)

    return result


def get_bond_total_return(bond_code, start_date, end_date):
    """
    得到从start_date到end_date的时间区间内，每支债券的区间收益
    取这段时间的日行情，向后填充，再向前填充

    :param str or list bond_code: 需要的债券代码列表，债券数据太多，必须传入该参数
    :param str or int or datetime.date or datetime.datetime or int start_date: 开始时间，支持带杠字符串、不带杠字符串、datetime、date、int格式
    :param str or int or datetime.date or datetime.datetime or int end_date: 结束时间，支持带杠字符串、不带杠字符串、datetime、date、int格式
    :return: pd.DataFrame
    """
    conn = db_conn.get_basic_data_conn()

    if isinstance(bond_code, str):
        bond_code = [bond_code]
    bond_where_condition = get_lst_condition(bond_code, 'bond_code')

    start_date = time_tool.format_date_str(start_date, '%Y%m%d')
    end_date = time_tool.format_date_str(end_date, '%Y%m%d')

    bond_hq = my_mysql.read_v2(select=['bond_code', 'end_date as tj', 'deeval_fullprice as close'], sfrom=BOND_HQ_TABLE,
                               where=[f"end_date >= '{start_date}'", f"end_date <= '{end_date}'", bond_where_condition],
                               conn=conn)
    if bond_hq.empty:
        bond_ret = pd.DataFrame(columns=['bond_code', 'ret'])
    else:
        bond_hq = bond_hq.drop_duplicates(subset=['bond_code', 'tj'])
        bond_hq = bond_hq.set_index(['bond_code', 'tj'])['close'].unstack()
        bond_hq = bond_hq.ffill(axis=1)
        bond_hq = bond_hq.bfill(axis=1)
        bond_hq['ret'] = bond_hq[bond_hq.columns[-1]] / bond_hq[bond_hq.columns[0]] - 1
        bond_ret = bond_hq['ret'].reset_index()

    # 可转债行情表中没有bond_code字段，需要得到bond_code和secinner_code的对应关系
    conv_bond_code = my_mysql.read_v2(select=['secinner_code', 'bond_code'], sfrom=BOND_INFO_TABLE,
                                      where=bond_where_condition, conn=conn)
    if conv_bond_code.empty:
        conv_bond_ret = pd.DataFrame(columns=['bond_code', 'ret'])
    else:
        secinner_code_where_condition = get_lst_condition(conv_bond_code['secinner_code'].tolist(), 'secinner_code')
        conv_bond_hq = my_mysql.read_v2(select=['secinner_code', 'trade_date as tj', 'close'],
                                        sfrom=CONVERTIBLE_BOND_HQ_TABLE,
                                        where=[f"trade_date >= '{start_date}'", f"trade_date <= '{end_date}'",
                                               secinner_code_where_condition],
                                        conn=conn)
        if conv_bond_hq.empty:
            conv_bond_ret = pd.DataFrame()
        else:
            conv_bond_hq = conv_bond_hq.set_index(['secinner_code', 'tj'])['close'].unstack()
            conv_bond_hq = conv_bond_hq.ffill(axis=1)
            conv_bond_hq = conv_bond_hq.bfill(axis=1)
            conv_bond_hq['ret'] = (conv_bond_hq[conv_bond_hq.columns[-1]] /
                                   conv_bond_hq[conv_bond_hq.columns[0]] - 1)
            conv_bond_ret = conv_bond_hq['ret'].reset_index()
            conv_bond_ret = conv_bond_ret.merge(conv_bond_code, how='left', on='secinner_code')[['bond_code', 'ret']]

    bond_ret = pd.concat([bond_ret, conv_bond_ret], ignore_index=True, sort=False)

    conn.close()

    return bond_ret


def get_bond_daily_return(bond_code, start_date, end_date):
    """
    得到从start_date到end_date的时间区间内，债券每天的涨跌幅
    取这段时间的日行情，向后填充，再向前填充

    :param str or list bond_code: 需要的债券代码列表，债券数据太多，必须传入该参数
    :param str or int or datetime.date or datetime.datetime or int start_date: 开始时间，支持带杠字符串、不带杠字符串、datetime、date、int格式
    :param str or int or datetime.date or datetime.datetime or int end_date: 结束时间，支持带杠字符串、不带杠字符串、datetime、date、int格式
    :return: pd.DataFrame
    """
    conn = db_conn.get_basic_data_conn()

    if isinstance(bond_code, str):
        bond_code = [bond_code]
    bond_where_condition = get_lst_condition(bond_code, 'bond_code')

    start_date = time_tool.format_date_str(start_date, '%Y%m%d')
    end_date = time_tool.format_date_str(end_date, '%Y%m%d')

    bond_hq = my_mysql.read_v2(select=['bond_code', 'end_date', 'deeval_fullprice as close'], sfrom=BOND_HQ_TABLE,
                               where=[f"end_date >= '{start_date}'", f"end_date <= '{end_date}'", bond_where_condition],
                               conn=conn)
    if bond_hq.empty:
        bond_hq = pd.DataFrame(columns=['bond_code', 'end_date', 'close'])
    else:
        bond_hq = bond_hq.drop_duplicates(subset=['bond_code', 'end_date'])
        bond_hq = bond_hq.set_index(['bond_code', 'end_date'])['close'].unstack()
        bond_hq = bond_hq.ffill(axis=1)
        bond_hq = bond_hq.bfill(axis=1)
        bond_hq = bond_hq.stack().rename('close').reset_index()

    # 可转债行情表中没有bond_code字段，需要得到bond_code和secinner_code的对应关系
    conv_bond_code = my_mysql.read_v2(select=['secinner_code', 'bond_code'], sfrom=BOND_INFO_TABLE,
                                      where=bond_where_condition, conn=conn)
    if conv_bond_code.empty:
        conv_bond_hq = pd.DataFrame(columns=['bond_code', 'end_date', 'close'])
    else:
        secinner_code_where_condition = get_lst_condition(conv_bond_code['secinner_code'].tolist(), 'secinner_code')
        conv_bond_hq = my_mysql.read_v2(select=['secinner_code', 'trade_date as end_date', 'close'],
                                        sfrom=CONVERTIBLE_BOND_HQ_TABLE,
                                        where=[f"trade_date >= '{start_date}'", f"trade_date <= '{end_date}'",
                                               secinner_code_where_condition],
                                        conn=conn)
        if conv_bond_hq.empty:
            conv_bond_hq = pd.DataFrame(columns=['bond_code', 'end_date', 'close'])
        else:
            conv_bond_hq = conv_bond_hq.set_index(['secinner_code', 'end_date'])['close'].unstack()
            conv_bond_hq = conv_bond_hq.ffill(axis=1)
            conv_bond_hq = conv_bond_hq.bfill(axis=1)
            conv_bond_hq = conv_bond_hq.stack().rename('close').reset_index().merge(conv_bond_code, how='left',
                                                                                    on='secinner_code')
            conv_bond_hq = conv_bond_hq.drop(columns=['secinner_code'])
    result = pd.concat([bond_hq, conv_bond_hq])
    result['end_date'] = result['end_date'].apply(time_tool.format_date_str)
    conn.close()

    return result


if __name__ == '__main__':
    aaa = get_bond_total_return('011104002', '2012-01-01', '2012-03-01')
    bbb = get_bond_daily_return('011104002', '2012-01-01', '2012-03-01')
    ccc = get_bond_quote('011104002', '2012-01-01', '2012-03-01')
